from gensim.models import Word2Vec
from sklearn.cluster import KMeans
from sklearn.decomposition import PCA
import numpy as np
import pandas as pd
import dash
from dash import dcc, html
from dash.dependencies import Input, Output
import plotly.express as px
import os
import time


# 定义版本
version = 31

# 加载已经训练好的 Word2Vec 模型
model = Word2Vec.load(f'./models/word2vec-tinymodel-v{version}.model')

# 获取词向量
word_vectors = []
word_list = list(model.wv.index_to_key)
for word in word_list:
    word_vectors.append(model.wv[word])

# 转换为 NumPy 数组
word_vectors = np.array(word_vectors)

# 使用 KMeans 聚类
n_clusters = 100 # 设置聚类数
kmeans = KMeans(n_clusters=n_clusters, random_state=0)
kmeans.fit(word_vectors)

# 获取聚类结果
labels = kmeans.labels_

# 使用 PCA 将 300 维降到 2 维
pca = PCA(n_components=2)
word_vectors_2d = pca.fit_transform(word_vectors)

# 将结果转换为 DataFrame 以便于可视化
df = pd.DataFrame(word_vectors_2d, columns=['x', 'y'])
df['word'] = word_list
df['cluster'] = labels

# 导出为 CSV 文件
output_csv_path = f'./kmeans_clusters_v{version}.csv'
df.to_csv(output_csv_path, index=False)
print(f"聚类结果已保存到 {output_csv_path}")

# 导出为 CSV 文件
output_csv_path = f'./kmeans_clusters_v{version}.csv'
df.to_csv(output_csv_path, index=False)
print(f"聚类结果已保存到 {output_csv_path}")

# 创建 Dash 应用
app = dash.Dash(__name__)

app.layout = html.Div([
    dcc.Dropdown(
        id='cluster-dropdown',
        options=[{'label': f'Cluster {i}', 'value': i} for i in range(n_clusters)],
        multi=True,
        value=[0],  # 默认选择一个聚类
        placeholder='Select cluster(s)'
    ),
    dcc.Graph(id='scatter-plot')
])

@app.callback(
    Output('scatter-plot', 'figure'),
    Input('cluster-dropdown', 'value')
)
def update_figure(selected_clusters):
    if not selected_clusters:
        filtered_df = df
    else:
        filtered_df = df[df['cluster'].isin(selected_clusters)]
    
    fig = px.scatter(filtered_df, x='x', y='y', color='cluster', text='word',
                     title='KMeans Clustering of Word Vectors',
                     labels={'x': 'Dimension 1', 'y': 'Dimension 2'},
                     color_continuous_scale=px.colors.sequential.Viridis)
    fig.update_traces(textposition='top center')
    fig.update_layout(width=1500, height=700)
    creation_time = os.path.getctime(f'./models/word2vec-tinymodel-v{version}.model')
    creation_time_readable = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(creation_time))
    fig.add_annotation(
        text=(
            f"模型版本 v{version} "
            f"训练时间：{creation_time_readable} "
            f"cluster: {n_clusters}<br>"
            f"v_size: {model.vector_size}, "
            f"window: {model.window}, "
            f"min_count: {model.min_count}, "
            f"sg: {model.sg}, "
            f"epochs: {model.epochs}, "
            f"negative: {model.negative}, "
            f"sample: {model.sample}, "
            f"words: {len(model.wv.index_to_key)}"
        ),
        xref="paper",  
        yref="paper",
        x=0.5,  
        y=-0.15,  
        showarrow=False,  
        font=dict(size=12),
    )
    
    return fig

if __name__ == '__main__':
    app.run_server(debug=True)
